package sspatcher;

import sspatcher.entity.PatchEntity;
import sspatcher.entity.PatchRuleEntity;

import java.io.*;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.ProtectionDomain;
import java.util.*;
import java.util.logging.FileHandler;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;

// java agent
public class PatcherAgent {
    public static final Logger logger = Logger.getLogger(PatcherAgent.class.getName());

    public static void premain(String agentArgs, Instrumentation inst) {
        Map<String, String> runConfig;
        try {
            runConfig = Utilities.readIniFromFile(findPatcherRunINI());
            checkPatcherRunININeededKeys(runConfig);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        File rootFolder = Paths.get(runConfig.get("PATCHER_ROOT")).toAbsolutePath().normalize().toFile();
        if (!rootFolder.exists() && !rootFolder.mkdir()) {
            throw new RuntimeException("cannot make directory:" + rootFolder);
        }
        File modsFolder = Paths.get(rootFolder.getPath(), runConfig.get("MODS_DIR")).toFile();
        if (!modsFolder.exists() && !modsFolder.mkdir()) {
            throw new RuntimeException("cannot make directory:" + modsFolder);
        }

        try {   // setup logger
            logger.addHandler(new FileHandler(Paths.get(rootFolder.getPath(), "patcher.log").toString(), false));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        List<String> patchList = readPatchList(Paths.get(rootFolder.getPath(), runConfig.get("PATCH_LIST")).toFile());
        if (patchList.isEmpty()) {
            logger.warning("no patch was enabled");
        }

        List<PatchEntity> patchEntities = patchList.stream().map(
                patchLine -> {
                    return loadPatchPackage(Paths.get(modsFolder.getPath(), patchLine).toString());
                }
        ).flatMap(Collection::stream).collect(Collectors.toList());

        PatcherTransformer pt = new PatcherTransformer();
        pt.addPatchEntities(patchEntities);
        inst.addTransformer(pt);
    }

    private static List<PatchEntity> loadPatchPackage(String iniPath) {
        Map<String, String> patchPackConfig;
        try {
            patchPackConfig = Utilities.readIniFromFile(new File(iniPath));
        } catch (IOException e) {
            logger.log(Level.SEVERE, "fail to read ini: " + iniPath, e);
            return new ArrayList<>();
        }

        String csvPath = Paths.get(new File(iniPath).getParent(), patchPackConfig.get("PATCH_TABLE")).toString();
        String modFolder = Paths.get(new File(iniPath).getParent(), patchPackConfig.get("RULE_FOLDER")).toString();
        try {
            return CSVHelper.readCsvFromFile(new File(csvPath), true).stream().map(
                    patchLine -> {
                        PatchEntity patchEntity = new PatchEntity();
                        patchEntity.setHash(patchLine.get(0));
                        patchEntity.setCsvPath(patchLine.get(1));
                        patchEntity.setTitle(patchLine.get(2));
                        patchEntity.setDescription(patchLine.get(3));
                        patchEntity.setRules(loadPatchRules(
                                Paths.get(modFolder, patchEntity.getCsvPath()).toString()
                        ));
                        return patchEntity;
                    }
            ).collect(Collectors.toList());
        } catch (Exception e) {
            logger.log(Level.SEVERE, "fail to read csv: " + csvPath, e);
            return new ArrayList<>();
        }
    }

    private static List<PatchRuleEntity> loadPatchRules(String csvPath) {
        try {
            return CSVHelper.readCsvFromFile(new File(csvPath), true).stream().map(
                    patchLine -> {
                        PatchRuleEntity rule = new PatchRuleEntity();
                        rule.setTypeEnum(PatchRuleTypeEnum.fromString(patchLine.get(0)));
                        rule.setIndex(Utilities.parseIntAutoHex(patchLine.get(1)));
                        rule.setData(patchLine.get(2));
                        rule.setExtern(patchLine.get(3));
                        return rule;
                    }
            ).collect(Collectors.toList());
        } catch (Exception e) {
            logger.log(Level.SEVERE, "fail to read csv: " + csvPath, e);
            return new ArrayList<>();
        }
    }

    private static void checkPatcherRunININeededKeys(Map<String, String> runConfig) {
        final String[] keys = new String[]{"PATCHER_ROOT", "MODS_DIR", "PATCH_LIST"};
        for (String key : keys) {
            if (!runConfig.containsKey(key)) {
                throw new RuntimeException("in PATCHER_RUN.ini, missing required key: " + key);
            }
        }
    }

    private static File findPatcherRunINI() {
        File runFileInPatcher = new File("./patcher/PATCHER_RUN.ini");
        if (runFileInPatcher.exists()) return runFileInPatcher;
        File runFileInRoot = new File("./PATCHER_RUN.ini");
        if (runFileInRoot.exists()) return runFileInRoot;
        throw new RuntimeException("cannot find config file PATCHER_RUN.ini");
    }

    public static List<String> readPatchList(File patchListFile) {
        try {
            return Utilities.readListFromFile(patchListFile);
        }
        catch (IOException e) {
            logger.log(Level.SEVERE, "cannot read patch list: " + patchListFile, e);
            return new ArrayList<>();
        }
    }

    static class PatcherTransformer implements ClassFileTransformer {
        private int activePatchCnt = 0;
        private int matchedPatchCnt = 0;

        private final Map<String, List<PatchRuleEntity>> mapHashToPatches = new HashMap<>();
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
            String hex = Utilities.calcMD5LowerCast(classfileBuffer);
            List<PatchRuleEntity> patches = mapHashToPatches.getOrDefault(hex, null);
            if (null == patches) {
                return classfileBuffer;
            }

            logger.info(String.format("match (%d/%d)", ++matchedPatchCnt, activePatchCnt));
            logger.info("applying patch to matched hash: " + hex);
            try {
                byte[] bytecode = ASMHelper.applyPatchToBytecode(classfileBuffer, patches);
//                logger.info(Arrays.toString(bytecode));
                return bytecode;
            }
            catch (Exception e) {
                logger.log(Level.SEVERE, "failed to apply patch: " + hex, e);
            }

            return classfileBuffer;
        }

        public void addPatchEntity(PatchEntity patchEntity) {
            String h = patchEntity.getHash().toLowerCase();
            if (mapHashToPatches.containsKey(h)) {
                logger.warning("skipping patch because of duplicated patch hash: " + h);
                return;
            }
            activePatchCnt++;
            mapHashToPatches.put(h, patchEntity.getRules());
        }

        public void addPatchEntities(List<PatchEntity> patchEntities) {
            for (PatchEntity pe : patchEntities) {
                addPatchEntity(pe);
            }
        }
    }
}